// matrix.go
package matrix

import (
	"fmt"
	"log"
	"wgmat/compute"
	"wgmat/gpucore"
)

// Matrix 表示一个 GPU 上的矩阵，支持链式操作。
// 它内部持有一个 DeferredCompute 和一个 MatrixRef。
// 所有操作都是延迟的，直到调用 Fetch()。
type Matrix struct {
	dc  *compute.DeferredCompute
	ref compute.MatrixRef
	err error // 错误传播（类似 errgroup）
}

// NewMatrixFromData 从 CPU 数据创建一个 GPU 矩阵（立即上传）
func NewMatrixFromData(ctx gpucore.GPUContext, data []float32, rows, cols uint32) (*Matrix, error) {
	if len(data) != int(rows*cols) {
		return nil, fmt.Errorf("data length %d != rows*cols %d", len(data), rows*cols)
	}
	buf, err := ctx.CreateBuffer(data, gpucore.BufferUsageStorage|gpucore.BufferUsageCopySrc)
	if err != nil {
		return nil, fmt.Errorf("upload matrix: %w", err)
	}
	ref := compute.NewOwnedMatrixRef(buf, rows, cols)
	dc, err := compute.NewDeferredCompute(ctx)
	if err != nil {
		buf.Release()
		return nil, fmt.Errorf("create deferred compute: %w", err)
	}
	return &Matrix{dc: dc, ref: ref}, nil
}

// NewMatrixFromRef 从已有 MatrixRef 创建（用于中间结果）
func NewMatrixFromRef(dc *compute.DeferredCompute, ref compute.MatrixRef) *Matrix {
	return &Matrix{dc: dc, ref: ref}
}

// Add 返回 this + other
func (m *Matrix) Add(other *Matrix) *Matrix {
	if m.err != nil {
		return m
	}
	if other.err != nil {
		return &Matrix{err: other.err}
	}
	if m.dc.Device() != other.dc.Device() {
		return &Matrix{err: fmt.Errorf("matrices from different GPU contexts")}
	}
	// 注意：我们复用 m.dc，但 other 的操作必须已编码进同一个 encoder！
	// 为简化，要求所有操作在同一个 DeferredCompute 上构建。
	// 因此：**禁止跨 DeferredCompute 操作**。
	// 解决方案：强制所有中间 Matrix 共享同一个 dc。
	ref, err := m.dc.Add(m.ref, other.ref)
	if err != nil {
		return &Matrix{err: err}
	}
	return NewMatrixFromRef(m.dc, ref)
}

// Mul 返回 this * other（矩阵乘法）
func (m *Matrix) Mul(other *Matrix) *Matrix {
	if m.err != nil {
		return m
	}
	if other.err != nil {
		return &Matrix{err: other.err}
	}
	if m.ref.Cols != other.ref.Rows {
		return &Matrix{err: fmt.Errorf("matmul dim mismatch: %dx%d * %dx%d", m.ref.Rows, m.ref.Cols, other.ref.Rows, other.ref.Cols)}
	}
	params := compute.GemmParams{M: m.ref.Rows, N: other.ref.Cols, K: m.ref.Cols}
	ref, err := m.dc.GEMM(m.ref, other.ref, params)
	if err != nil {
		return &Matrix{err: err}
	}
	return NewMatrixFromRef(m.dc, ref)
}

// LUFactorize 返回 (LU, pivot)，其中 LU 是 in-place 分解结果
func (m *Matrix) LUFactorize() (lu *Matrix, pivotBuf gpucore.GPUBuffer, err error) {
	if m.err != nil {
		return nil, nil, m.err
	}
	luRef, pivBuf, err := m.dc.LUFactorize(m.ref)
	if err != nil {
		return nil, nil, err
	}
	return NewMatrixFromRef(m.dc, luRef), pivBuf, nil
}

// Solve 解 Ax = b，其中 A 是 this（已 LU 分解），b 是列向量
func (m *Matrix) Solve(b *Matrix) *Matrix {
	if m.err != nil {
		return m
	}
	if b.err != nil {
		return &Matrix{err: b.err}
	}
	// pivotBuf 从 lu 矩阵的 ref 中获取
	xRef, err := m.dc.SolveGPU(m.ref, b.ref, m.ref.PivotBuf)
	if err != nil {
		return &Matrix{err: err}
	}
	return NewMatrixFromRef(m.dc, xRef)
}

func (m *Matrix) Fetch() ([]float32, error) {
	if m.err != nil {
		return nil, m.err
	}

	log.Printf("Fetch: 提交GPU操作...")
	if err := m.dc.Submit(); err != nil {
		return nil, fmt.Errorf("submit: %w", err)
	}

	// 在 Submit 后检查奇异性
	if m.ref.PivotBuf != nil {
		log.Printf("Fetch: 检查奇异性...")
		if err := m.dc.CheckSingularity(m.ref.PivotBuf); err != nil {
			return nil, fmt.Errorf("singular matrix detected: %w", err)
		}
	}

	log.Printf("Fetch: 读取结果...")
	data, err := m.ref.Buffer.Read()
	if err != nil {
		return nil, fmt.Errorf("read result: %w", err)
	}
	return data, nil
}

// matrix.go

// CheckSingularityAfterLU 在 LUFactorize 后检查矩阵是否奇异
// 这会提交当前的延迟操作并检查 pivot 缓冲区
// CheckSingularityAfterLU 在 LUFactorize 后检查矩阵是否奇异
func (m *Matrix) CheckSingularityAfterLU() error {
	if m.err != nil {
		return m.err
	}
	if m.ref.PivotBuf == nil {
		return fmt.Errorf("no pivot buffer available")
	}

	if err := m.dc.Submit(); err != nil {
		return fmt.Errorf("submit: %w", err)
	}

	return m.dc.CheckSingularity(m.ref.PivotBuf)
}

// FetchPivot 安全读取 pivot 缓冲区（必须在 GPU 操作完成后调用）
func FetchPivot(pivBuf gpucore.GPUBuffer) ([]uint32, error) {
	data, err := pivBuf.ReadAsUint32()
	if err != nil {
		return nil, err
	}

	// 检查奇异性标记
	if len(data) > 0 && data[0] == 0xFFFFFFFF {
		return nil, fmt.Errorf("matrix is singular")
	}

	log.Printf("FetchPivot: pivot = %v", data)
	return data, nil
}

// Close 释放资源（可选，Fetch 后通常不再需要）
func (m *Matrix) Close() {
	m.ref.Close()
	m.dc.Close()
}

// Err returns the error, if any operation failed.
func (m *Matrix) Err() error {
	return m.err
}
